import torch
import networkx as nx
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE


def visualize_graph(G, color):
    plt.figure(figsize=(7, 7))
    plt.xticks([])
    plt.yticks([])
    nx.draw_networkx(G, pos=nx.spring_layout(G, seed=42), with_labels=False, node_color=color, cmap="Set2")
    # plt.show()
    # 保存图像为 PNG 格式
    plt.savefig('节点关系图.png')


def visualize_embedding(h, color, epoch=None, loss=None):
    plt.figure(figsize=(7, 7))
    plt.xticks([])
    plt.yticks([])
    h = h.detach().cpu().numpy()
    plt.scatter(h[:, 0], h[:, 1], s=140, c=color, cmap="Set2")
    if epoch is not None and loss is not None:
        plt.xlabel("Epoch {},Loss {}".format(epoch, loss.item()), fontsize=16)
        plt.savefig('隐藏特征第{}轮信息.png'.format(epoch))
    else:
        plt.savefig('隐藏特征初始信息.png')


def visualize(h, color):
    z = TSNE(n_components=2).fit_transform(h.detach().cpu().numpy())
    plt.figure(figsize=(10, 10))
    plt.xticks([])
    plt.yticks([])
    plt.scatter(z[:, 0], z[:, 1], s=70, c=color, cmap="Set2")
    # 保存图像为 PNG 格式
    plt.savefig('节点关系图.png')
